-- 创建一个新的协议
tls_protocol = Proto("hypertls", "Kernel TLS Protocol")

local record_type_values = {
  [22] = "Handshake",
  [23] = "Appliaction"
}

local handshake_version_values = {
  [0x301] = "TLS1.0",
  [0x302] = "TLS1.2",
  [0x303] = "TLS1.3",
}

local handshake_type_values = {
  [1] = "Client Hello",
  [2] = "Server Hello",
  [11] = "Certificate",
  [14] = "Finished",

}

local info_values = {
  [1] = "Client Hello",
  [2] = "Server Hello, Certificate, Server Finished, Signature",
  [11] = "Certificate, Client Finished, Signature",
  [14] = "Finished",

}

record_type_pf = ProtoField.uint8(
  "tls.record_type", "Content Type", base.DEC, record_type_values)
finish_type_pf = ProtoField.uint8(
  "tls.finish", "Finished", base.DEC)
version_pf = ProtoField.uint16(
  "tls.version", "Version", base.HEX, handshake_version_values)

length_pf = ProtoField.uint16(
  "tls.length", "Length", base.DEC)

handshake_type_pf = ProtoField.uint8(
  "tls.hanshake_type", "hanshake Type", base.DEC, handshake_type_values)

random_pf = ProtoField.bytes(
  "tls.random", "Random", base.NONE)
curve_id_pf = ProtoField.uint16(
  "tls.curveid", "Curve_id", base.HEX)
keyshare_pf = ProtoField.bytes(
  "tls.keyshare", "keyshare", base.NONE)
cert_pf = ProtoField.bytes(
  "tls.cert", "Cert:", base.NONE)
sign_pf = ProtoField.bytes(
  "tls.sign", "Sign:", base.NONE)
  hmac_pf = ProtoField.bytes(
    "tls.hmac", "Hmac:", base.NONE)
-- 定义协议的字段
tls_protocol.fields = {
  record_type_pf, version_pf, length_pf, handshake_type_pf, random_pf, curve_id_pf, keyshare_pf, cert_pf, finish_type_pf,
  sign_pf, hmac_pf
}

-- 解析TLS记录的函数
function tls_protocol.dissector(buffer, pinfo, tree)
  -- 设置协议名称
  pinfo.cols.protocol = "HyperTLS"
  local type_buf = buffer(0, 1)
  local version_buf = buffer(1, 2)
  local length_buf = buffer(3, 2)
  local handshake_type_buf = buffer(5, 1)
  local handshake_length_buf = buffer(7, 2)
  local handshake_version_buf = buffer(9, 2)
  local random_buf = buffer(11, 32)
  local curve_id = buffer(43, 2)
  local keyshare_len_buf = buffer(45, 2)
  local keyshare_buf = buffer(47, 32)


  -- 创建一个树状结构来展示TLS记录
  local subtree = tree:add(tls_protocol, buffer())

  -- 解析record layer
  local record_layer = subtree:add("Record Layer")

  record_layer:add(record_type_pf, type_buf)
  record_layer:add(version_pf, version_buf)
  record_layer:add(length_pf, length_buf)

  -- 解析handshake layer
  local handshake_type = handshake_type_buf:uint()
  local handshake_typer_str = handshake_type_values[handshake_type]
  local handshake_layer = record_layer:add("Handshake Protocol: " .. handshake_typer_str)
  pinfo.cols.info = info_values[handshake_type]
  handshake_layer:add(handshake_type_pf, handshake_type_buf)

  if handshake_type == 0xb then
    local version_buf = buffer(7, 2)
    local length_buf = buffer(9, 2)
    handshake_layer:add(version_pf, version_buf)
    handshake_layer:add(length_pf, length_buf)
    local cert = handshake_layer:add("Certificate")
    local cert_length = length_buf:uint()
    cert:add(cert_pf, buffer(11, cert_length))
    local record_layer = subtree:add("Record Layer")
    record_layer:add(record_type_pf, buffer(11 + cert_length, 1))
    record_layer:add(version_pf, buffer(12 + cert_length, 2))
    record_layer:add(length_pf, buffer(14 + cert_length, 2))
    local handshake_type_buf = buffer(16 + cert_length, 1)
    handshake_type = handshake_type_buf:uint()
    handshake_typer_str = handshake_type_values[handshake_type]
    local handshake = record_layer:add("Handshake Protocol: " .. handshake_typer_str)
    handshake:add(hmac_pf, buffer(20 + cert_length, 32))
    record_layer = subtree:add("Record Layer")
    record_layer:add(record_type_pf, buffer(52 + cert_length, 1))
    record_layer:add(version_pf, buffer(53 + cert_length, 2))
    record_layer:add(length_pf, buffer(55 + cert_length, 2))
    local sign = record_layer:add("Signature")
    sign:add(sign_pf, buffer(57 + cert_length, 512))
    return
  end
  handshake_layer:add(version_pf, handshake_version_buf)
  handshake_layer:add(length_pf, handshake_length_buf)
  handshake_layer:add(random_pf, random_buf)

  local keyshare = handshake_layer:add("keyshare")
  keyshare:add(curve_id_pf, curve_id):append_text(" (" .. "x25519" .. ")")
  keyshare:add(length_pf, keyshare_len_buf)
  keyshare:add(keyshare_pf, keyshare_buf)

  if handshake_type == 2 then
    local server_hello_type_buf = buffer(79, 1)
    local server_hello_version_buf = buffer(80, 2)
    local server_hello_length_buf = buffer(82, 2)
    local handshake_type_buf = buffer(84, 1)
    local handshake_version_buf = buffer(86, 2)
    local handshake_length_buf = buffer(88, 2)
    record_layer = subtree:add("Record Layer")
    record_layer:add(record_type_pf, server_hello_type_buf)
    record_layer:add(version_pf, server_hello_version_buf)
    record_layer:add(length_pf, server_hello_length_buf)
    local handshake_type = handshake_type_buf:uint()
    local handshake_typer_str = handshake_type_values[handshake_type]
    local handshake_layer = record_layer:add("Handshake Protocol: " .. handshake_typer_str)
    handshake_layer:add(handshake_type_pf, handshake_type_buf)
    handshake_layer:add(version_pf, handshake_version_buf)
    handshake_layer:add(length_pf, handshake_length_buf)
    local cert = handshake_layer:add("Certificate")
    local cert_length = handshake_length_buf:uint()
    cert:add(cert_pf, buffer(90, cert_length))

    local record_layer = subtree:add("Record Layer")
    record_layer:add(record_type_pf, buffer(90 + cert_length, 1))
    record_layer:add(version_pf, buffer(91 + cert_length, 2))
    record_layer:add(length_pf, buffer(93 + cert_length, 2))
    handshake_type_buf = buffer(95 + cert_length, 1)
    handshake_type = handshake_type_buf:uint()
    handshake_typer_str = handshake_type_values[handshake_type]
    handshake = record_layer:add("Handshake Protocol: " .. handshake_typer_str)
    handshake:add(finish_type_pf, buffer(95 + cert_length, 4))
    local record_layer = subtree:add("Record Layer")
    record_layer:add(record_type_pf, buffer(99 + cert_length, 1))
    record_layer:add(version_pf, buffer(100 + cert_length, 2))
    record_layer:add(length_pf, buffer(102 + cert_length, 2))
    local sign = record_layer:add("Signature")
    sign:add(sign_pf, buffer(104 + cert_length, 512))
  end
end

-- 将协议应用到Wireshark的TCP端口9999
local tcp_port = DissectorTable.get("tcp.port")
tcp_port:add(9999, tls_protocol)
